package top.ieei.demo.tftp;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.SocketException;
import java.net.UnknownHostException;

/**
 * TFTP客户端<br>
 * 
 * @author xiecongshu
 */
public class TftpClient
{

    /**
     * 默认TFTP
     */
    public static final int DEFAULT_PORT = 69;

    /**
     * TFTP服务IP
     */
    public static final String DEFAULT_HOSTNAME = "127.0.0.1";

    /**
     * 最大包大小
     */
    private static final int MAX_PACKAGE_SIZE = 512;

    /**
     * the name of the host
     */
    private String hostName = null;

    /**
     * TFTP服务端口
     */
    private int port = DEFAULT_PORT;

    /**
     * 超时时间.
     */
    private int timeout;

    /**
     * 传输文件大小
     */
    private int tsize;

    private boolean stuffedLink = true;

    /**
     * 
     * 构造函数<br>
     * 
     * @author xiecongshu
     * @since 1.0, 2013-5-7
     * @param hostName
     */
    public TftpClient(String hostName)
    {
        this.setHostName(hostName);
    }

    /**
     * 
     * 构造函数<br>
     * 
     * @author xiecongshu
     * @since 1.0, 2013-5-7
     * @param hostName
     * @param port
     */
    public TftpClient(String hostName, int port)
    {
        this.setHostName(hostName);
        this.port = port;
    }

    /**
     * @return Returns the hostName.
     */
    public String getHostName()
    {
        return hostName;
    }

    /**
     * @param hostName The hostName to set.
     */
    public void setHostName(String hostName)
    {
        this.hostName = hostName;
    }

    /**
     * @return Returns the port.
     */
    public int getPort()
    {
        return port;
    }

    /**
     * @param port The port to set.
     */
    public void setPort(int port)
    {
        this.port = port;
    }

    public ReadRequest initialiseDownload(String fileName, int optionTimeout, int optionTransferSize)
            throws UnknownHostException, TftpException
    {
        // The following code sets up a read request to the server, providing
        // the server
        // with the file name that it wants the data for.
        ReadRequest rrq = new ReadRequest();
        rrq.setFileName(fileName);
        if (optionTimeout != 0)
        {
            rrq.setTimeout(optionTimeout);
        }
        if (optionTransferSize != 0)
        {
            rrq.setTransferSize(optionTransferSize);
        }

        rrq.setMode(TftpRequest.OCTET_MODE);
        rrq.setAddress(InetAddress.getByName(this.getHostName()));
        rrq.setPort(this.port);

        return rrq;
    }

    public WriteRequest initialiseUpload(String fileName, int optionTimeout, int optionTransferSize)
            throws UnknownHostException, TftpException
    {
        // The following code sets up a write request to the server.
        WriteRequest wrq = new WriteRequest();
        wrq.setFileName(fileName);
        if (optionTimeout != 0)
        {
            wrq.setTimeout(optionTimeout);
        }
        if (optionTransferSize != 0)
        {
            wrq.setTransferSize(optionTransferSize);
        }

        wrq.setMode(TftpRequest.OCTET_MODE);
        wrq.setAddress(InetAddress.getByName(getHostName()));
        wrq.setPort(this.port);

        return wrq;
    }

    public boolean download(ReadRequest rrq, OutputStream os) throws SocketException, TftpException, IOException
    {
        // create a TFTP Socket
        TftpSocket tftpSock = new TftpSocket(5);

        int sequenceNumber = 1;

        byte[] dummyByteArray = new byte[1]; // used in constructor for
                                             // DATA....doesnt do anything
        TftpData receive = new TftpData(sequenceNumber, dummyByteArray);
        TftpAck surprisePacket = this.sendRequest(tftpSock, rrq, receive);

        if (surprisePacket == null)
        {
            return false;
        }

        if (surprisePacket instanceof TftpOptionAck)
        {
            // we received some extra's
            TftpOptionAck oack = (TftpOptionAck) surprisePacket;
            tsize = oack.getTransferSize();
            timeout = oack.getTimeout();

            TftpAck ack = new TftpAck(0);
            ack.setPort(surprisePacket.getPort());
            ack.setAddress(surprisePacket.getAddress());

            receive = (TftpData) transferData(tftpSock, ack, receive);
            if (receive == null)
            {
                // nothing returned from server...on acking the oack.
                System.out.println("Nothing returned from the server after ack on oack.");
                return false;
            }
        }
        else if (surprisePacket instanceof TftpData)
        {
            // too bad, now we need to work
            receive = (TftpData) surprisePacket;
        }

        // write this data to the output stream....
        os.write(receive.getData());

        // need to find the port and address the server has chosen to
        // communicate on and connect to it.
        int serverPort = receive.getPort();
        System.out.println("The server has chosen the following port as the communication port: " + serverPort);
        InetAddress serverAddress = rrq.getAddress();
        tftpSock.connect(serverAddress, serverPort);

        TftpAck ack = null;
        while (receive.getData().length >= MAX_PACKAGE_SIZE)
        {
            ack = new TftpAck(sequenceNumber++);
            receive = new TftpData(sequenceNumber, dummyByteArray);
            System.out.println("receiving block" + sequenceNumber);

            receive = (TftpData) transferData(tftpSock, ack, receive);

            if (receive == null)
            {
                // nothing returned from server....
                System.out.println("Nothing returned from the server after the transfer.");
                return false;
            }
            os.write(receive.getData());
        }
        // now that the last packet in the file has been sent, the client must
        // sent an
        // acknowledgement to confirm it has received the last package...or else
        // the server
        // tries to resend..and resend....etc
        System.out.println("send ack to say that we have received last message.");
        ack = new TftpAck(sequenceNumber);
        receive = (TftpData) transferData(tftpSock, ack, null);

        // ensure that the stream is closed.
        os.close();

        return true;
    }

    /**
     * 
     * (一句话功能简述)<br>
     * 
     * @author xiecongshu
     * @since 1.0, 2013-5-7
     * @param wrq
     * @param is
     * @return
     * @throws SocketException
     * @throws TftpException
     * @throws IOException
     */
    public boolean upload(WriteRequest wrq, InputStream is) throws SocketException, TftpException, IOException
    {
        System.out.println("[M] : upload: " + wrq.getFileName());
        // create a TFTP Socket
        TftpSocket tftpSock = new TftpSocket(5);

        int sequenceNumber = 0;

        TftpAck receive = new TftpAck(0);
        receive = (TftpAck) this.sendRequest(tftpSock, wrq, receive);

        // the server should send an ACK (where block number=0, need to
        // check...)
        if (receive == null)
        {
            // nothing returned from server...should throw an exception at this
            // point.
            System.out.println("Nothing returned from the server after the initial send.");
            return false;
        }
        if (receive.getBlockNum() != 0)
        {
            // nothing returned from server...should throw an exception at this
            // point.
            System.out.println("The server has sent an ACK with wrong block number.");
            return false;
        }

        if (receive instanceof TftpOptionAck)
        {
            // we received some extra's
            TftpOptionAck oack = (TftpOptionAck) receive;
            tsize = oack.getTransferSize();
            timeout = oack.getTimeout();
        }

        // need to find the port and address the server has chosen to
        // communicate on and connect to it.
        int serverPort = receive.getPort();
        System.out.println("The server has chosen the following port as the communication port: " + serverPort);
        InetAddress serverAddress = wrq.getAddress();
        tftpSock.connect(serverAddress, serverPort);

        // need to read the file and send to server....
        byte[] sendBytes = new byte[MAX_PACKAGE_SIZE];
        TftpData send = new TftpData();
        int returnValue = 0;

        while ((returnValue = is.read(sendBytes)) != -1)
        {
            System.out.println("sending packet number: " + sequenceNumber);

            // need to convert the byte array into correct TFTP format for the
            // DATA obejct
            send = new TftpData(++sequenceNumber, sendBytes, 0, returnValue);
            receive = new TftpAck(sequenceNumber);

            // now send to server, which in turn sends an acknowledgement
            receive = (TftpAck) transferData(tftpSock, send, receive);

        }
        // must remember to close the inputstream!
        is.close();

        return true;
    }

    /**
     * This method is responsible for the initial communication between the
     * client and the server .
     * 
     * @param tftpSock
     * @param frq
     * @param recv
     * @return
     */
    public TftpAck sendRequest(TftpSocket tftpSocket, TftpRequest send, TftpAck recv)
    {

        int retries = 0;
        int failCount = 0;

        // 发送数据包
        try
        {
            tftpSocket.write(send);
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }

        TftpPacket tftpPacket = null;
        boolean receiving = true;

        // 等待成功响应ACK
        while (receiving)
        {
            try
            {
                // 如果连接满了，设置两倍超时时间
                if (stuffedLink)
                {
                    tftpSocket.setSockTimeout(timeout * 2);
                }
                tftpPacket = tftpSocket.read();
            }
            catch (Exception e)
            {
                e.printStackTrace();
                return null;
            }

            // 如果没有收到响应
            if (tftpPacket == null)
            {
                if (retries++ > 5)
                {
                    // 重试次数太多，丢弃处理，发送错误响应给客户端
                    sendErrorPacket(tftpSocket, TftpError.ERROR_NOT_DEFINED, "Retransmit limit exceeded.");
                    return null;
                }
                else
                {
                    try
                    {
                        tftpSocket.write(send);

                        // 客户端数据处理管道可能已经充满，需要设置较长超时时间给客户端清理，并给出回应
                        stuffedLink = true;
                    }
                    catch (Exception e)
                    {
                        e.printStackTrace();
                    }
                    continue;
                }
            }

            // 收到出错的数据包，结束传输请求
            if (tftpPacket instanceof TftpError)
            {
                System.out.println(((TftpError) tftpPacket).getErrorMessage());
                return null;
            }

            // 收到预期的回复响应
            if ((tftpPacket instanceof TftpAck) && verifyAnswerPacket(recv, (TftpAck) tftpPacket))
            {
                try
                {
                    // 正确收到响应包 将超时时间改回默认
                    if (stuffedLink)
                    {
                        stuffedLink = false;
                        tftpSocket.setSockTimeout(timeout);
                    }
                }
                catch (SocketException e)
                {
                    e.printStackTrace();
                }
                return (TftpAck) tftpPacket;
            }

            // 收到一些垃圾包，数量过多时，放弃
            if (failCount++ > 5)
            {
                return null;
            }
        }
        return null;
    }

    /**
     * 验证接收的数据包是否和预期的一致
     * 
     * @param expecting 期望的ACK
     * @param received 接收到的ACK
     * @return true 如果匹配预期的结果
     */
    private boolean verifyAnswerPacket(TftpAck expectation, TftpAck received)
    {
        if (received == null)
        {
            return false;
        }

        if (expectation == null)
        {
            return false;
        }
        // 操作码不一致
        if (expectation.getOpCode() != received.getOpCode())
        {
            return false;
        }
        // 传输数据块编号不一致
        if (expectation.getBlockNum() != received.getBlockNum())
        {
            return false;
        }
        return true;
    }

    /**
     * 传输数据
     * 
     * @param send
     * @param recv
     * @return TftpAck
     */
    public TftpAck transferData(TftpSocket tftpSocket, TftpAck send, TftpAck recv)
    {
        int retries = 0;
        int failCount = 0;

        // 发送数据包
        try
        {
            tftpSocket.write(send);
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }

        TftpPacket tftpPacket = null;
        boolean receiving = true;

        // 等待成功响应ACK
        while (receiving)
        {
            try
            {
                // 如果连接满了，设置两倍超时时间
                if (stuffedLink)
                {
                    tftpSocket.setSockTimeout(timeout * 2);
                }
                tftpPacket = tftpSocket.read();
            }
            catch (Exception e)
            {
                e.printStackTrace();
                return null;
            }

            // 如果没有收到响应
            if (tftpPacket == null)
            {
                if (retries++ > 5)
                {
                    // 重试次数太多，丢弃处理，发送错误响应给客户端
                    sendErrorPacket(tftpSocket, TftpError.ERROR_NOT_DEFINED, "Retransmit limit exceeded.");
                    System.out.println(" Maximum retransmit count exceeded.");
                    return null;
                }
                else
                {
                    // 重发数据并等待
                    System.out.println("expected packet before time out, sending ACK/DATA again.");
                    try
                    {
                        tftpSocket.write(send);

                        // 客户端数据处理管道可能已经充满，需要设置较长超时时间给客户端清理，并给出回应
                        stuffedLink = true;
                    }
                    catch (Exception e)
                    {
                        e.printStackTrace();
                    }
                    continue;
                }
            }

            // 收到出错的数据包，结束传输请求
            if (tftpPacket instanceof TftpError)
            {
                System.out.println(((TftpError) tftpPacket).getErrorMessage());
                return null;
            }

            // 收到预期的回复响应
            if ((tftpPacket instanceof TftpAck) && verifyAnswerPacket(recv, (TftpAck) tftpPacket))
            {
                try
                {
                    // 正确收到响应包 将超时时间改回默认
                    if (stuffedLink)
                    {
                        stuffedLink = false;
                        tftpSocket.setSockTimeout(timeout);
                    }
                }
                catch (SocketException e)
                {
                    e.printStackTrace();
                }
                return (TftpAck) tftpPacket;
            }

            // 收到一些垃圾包，数量过多时，放弃
            if (failCount++ > 5)
            {
                return null;
            }
        }
        return null;
    }

    /**
     * 发送出错包回应客户端
     * 
     * @param errorCode
     * @param errorMsg void
     */
    public void sendErrorPacket(TftpSocket tftpSocket, int errorCode, String errorMsg)
    {
        try
        {
            TftpPacket tftpP = new TftpError(errorCode, errorMsg);
            tftpSocket.write(tftpP);
        }
        catch (Exception e)
        {
            e.printStackTrace();
            return;
        }
        finally
        {
        }
    }

    public int getOptionTimeout()
    {
        return timeout;
    }

    public int getOptionTransferSize()
    {
        return tsize;
    }

    /**
     * (一句话功能简述)<br>
     * 
     * @author xiecongshu
     * @since 1.0, 2013-5-7
     * @param args
     */
    public static void main(String[] args)
    {
        try
        {
            TftpClient client = new TftpClient("127.0.0.1");

            // invoke the download mechanism. This involves the device making a
            // Read
            // Request to the server.
            ReadRequest rrq = client.initialiseDownload("111.txt", 0, 0);

            // create an output stream to write the data to.
            FileOutputStream readFos = new FileOutputStream(new File("/temp/hello.txt"));
            client.download(rrq, readFos);

            // invoke the upload mechanism. This requires the device to make a
            // Write
            // Request to the server.
            WriteRequest wrq = client.initialiseUpload("", 0, 0);

            // create an output stream to write the data to.
            // need to send the data to the server...
            FileInputStream writeFis = new FileInputStream(new File("/temp/" + "sss.txt"));
            client.upload(wrq, writeFis);
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }
    }

}
